import argparse
from collections import defaultdict

def read_bed(input_file):
    sites = defaultdict(list)
    with open(input_file, 'r') as rf:
        for line in rf:
            words = line.strip().split("\t")
            sites[(words[0], words[5])].append(line)
    return sites

def filter(sites, outfilefh, non_outfilefh_hh, non_outfilefh_hl, non_outfilefh_lh, non_outfilefh_ll, kmer=21):
    kmer = kmer // 2
    for key in sites.keys():
        for i in range(len(sites[key])):
            iline = sites[key][i]
            iwords = iline.strip().split("\t")
            ipos = iwords[1]
            singleton_flag = 1
            value = int(iwords[10])
            if value != 100 and value != 0:
                continue
            for j in range(max(0, i - kmer), min(i + kmer, len(sites[key]))):
                if i == j:
                    continue
                jline = sites[key][j]
                jwords = jline.strip().split("\t")
                jpos = jwords[1]
                jvalue = int(jwords[10])
                if int(jpos) < int(ipos) + kmer and int(jpos) > int(ipos) - kmer:
                    if value == 100 and jvalue <= 10:
                        non_outfilefh_hl.append(jline)
                    elif value == 100 and jvalue > 10:
                        non_outfilefh_hh.append(jline)
                    elif value == 0 and jvalue >= 90:
                        non_outfilefh_lh.append(jline)
                    elif value == 0 and jvalue < 90:
                        non_outfilefh_ll.append(jline)
                    singleton_flag = 0
            if singleton_flag != 0:
                outfilefh.append(iline)

def write_to_file(output_file, singleton, hh, hl, lh, ll):
    with open(output_file, "w") as f:
        # Write singletons
        f.write("# Singleton results\n")
        for line in singleton:
            f.write(line)
        f.write("\n# HH group results\n")
        for line in hh:
            f.write(line)
        f.write("\n# HL group results\n")
        for line in hl:
            f.write(line)
        f.write("\n# LH group results\n")
        for line in lh:
            f.write(line)
        f.write("\n# LL group results\n")
        for line in ll:
            f.write(line)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/004_0111_001_R1_val_1_bismark_bt2_pe.deduplicated.bedGraph.gz.bismark.zero.cov.bed",
        help="input address.",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/home/xiaoyifu/data/HG002/Bisulfite/group_results.bisulfite.tsv",
        help="output file address.",
    )
    parser.add_argument(
        "--kmer",
        type=int,
        default=21,
        help="kmer size.",
    )
    return parser.parse_args()

if __name__ == '__main__':
    args = parse_args()
    sites = read_bed(args.input)
    
    # Define empty lists to store results
    singleton_results = []
    hh_results = []
    hl_results = []
    lh_results = []
    ll_results = []
    
    filter(sites, singleton_results, hh_results, hl_results, lh_results, ll_results, kmer=args.kmer)
    print('singleton_results:', len(singleton_results))
    print('hh_results:', len(hh_results))
    print('hl_results:', len(hl_results))
    print('lh_results:', len(lh_results))
    print('ll_results:', len(ll_results))   
    
    # Write all results to the same output file with separator lines
    write_to_file(args.output, singleton_results, hh_results, hl_results, lh_results, ll_results)
